"""生成计算与通信开销的 CSV（metrics.csv）。

测量：逐条验证 vs 聚合验证
输出列：N,method,avg_time_ms,comm_bytes
"""
import csv
import time
from cost import measure_computation_cost, measure_communication_cost
from kgc import kgc_partial_keygen
from user import user_full_key
from sign import sign_message, verify_signature
from aggregate import aggregate_signatures, verify_aggregate


def simulate_once(n: int):
    # 注册并生成密钥
    users = []
    for i in range(n):
        uid = f"vehicle-{i}".encode()
        partial = kgc_partial_keygen(uid)
        sk, pk = user_full_key(partial)
        users.append((sk, pk))

    messages = [b"speed=60" for _ in range(n)]
    signatures = [sign_message(sk, m) for (sk, _), m in zip(users, messages)]

    # 单条验证函数
    def verify_all():
        for (sk, pk), m, sig in zip(users, messages, signatures):
            # demo verify uses pk as sk
            verify_signature(pk, m, sig)

    # 聚合验证函数 (示例实现需要 individual signatures)
    agg_sig = aggregate_signatures(signatures)

    def verify_agg():
        verify_aggregate([pk for (_, pk) in users], messages, agg_sig, signatures=signatures)

    # 测量时间（单次平均，repeat=20）
    t_single = measure_computation_cost(verify_all, repeat=20)
    t_agg = measure_computation_cost(verify_agg, repeat=20)

    comm_single = measure_communication_cost(signatures)
    comm_agg = measure_communication_cost(agg_sig)

    return {
        "n": n,
        "single_ms": t_single,
        "agg_ms": t_agg,
        "single_bytes": comm_single,
        "agg_bytes": comm_agg,
    }


def main():
    sizes = [10, 50, 100, 200, 500]
    rows = []
    for n in sizes:
        print(f"Measuring n={n} ...")
        res = simulate_once(n)
        rows.append(res)

    # 写 CSV
    with open("metrics.csv", "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["N", "method", "avg_time_ms", "comm_bytes"])
        for r in rows:
            writer.writerow([r["n"], "single", f"{r['single_ms']:.6f}", r["single_bytes"]])
            writer.writerow([r["n"], "aggregate", f"{r['agg_ms']:.6f}", r["agg_bytes"]])

    print("metrics.csv generated")


if __name__ == "__main__":
    main()
